package raft

//
// this is an outline of the API that raft must expose to
// the service (or tester). see comments below for
// each of these functions for more details.
//
// rf = Make(...)
//   create a new Raft server.
// rf.Start(command interface{}) (index, term, isleader)
//   start agreement on a new log entry
// rf.GetState() (term, isLeader)
//   ask a Raft for its current term, and whether it thinks it is leader
// ApplyMsg
//   each time a new entry is committed to the log, each Raft peer
//   should send an ApplyMsg to the service (or tester)
//   in the same server.
//

import (
	"bytes"
	"math/rand"
	"sync"
	"time"
)
import "sync/atomic"
import "../labrpc"

import "../labgob"

const (
	follower = iota
	candidate
	leader
)

const (
	electionTimeout  = 300 * time.Millisecond
	heartBeatTimeout = 150 * time.Millisecond
)

//
// as each Raft peer becomes aware that successive log entries are
// committed, the peer should send an ApplyMsg to the service (or
// tester) on the same server, via the applyCh passed to Make(). set
// CommandValid to true to indicate that the ApplyMsg contains a newly
// committed log entry.
//
// in Lab 3 you'll want to send other kinds of messages (e.g.,
// snapshots) on the applyCh; at that point you can add fields to
// ApplyMsg, but set CommandValid to false for these other uses.
//
type ApplyMsg struct {
	CommandValid bool
	Command      interface{}
	CommandIndex int
}

type Entry struct {
	Index   int
	Term    int
	Command interface{}
}

type event struct {
	target interface{}
}

//
// A Go object implementing a single Raft peer.
//
type Raft struct {
	mu        sync.Mutex          // Lock to protect shared access to this peer's state
	peers     []*labrpc.ClientEnd // RPC end points of all peers
	persister *Persister          // Object to hold this peer's persisted state
	me        int                 // this peer's index into peers[]
	dead      int32

	// Your data here (2A, 2B, 2C).
	// Look at the paper's Figure 2 for a description of what
	// state a Raft server must maintain.

	// Persistent state
	currentTerm int
	votedFor    int
	log         []Entry
	// Volatile state
	commitIndex int
	lastApplied int
	state       int
	applyCh     chan ApplyMsg
	//leader state
	nextIndex  []int
	matchIndex []int

	stopped chan bool
	e       chan *event
}

func (rf *Raft) withLock(process func()) {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	process()
}

// return currentTerm and whether this server
// believes it is the leader.
func (rf *Raft) GetState() (int, bool) {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	// Your code here (2A).
	return rf.currentTerm, rf.state == leader
}

func (rf *Raft) State() int {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	return rf.state
}

func (rf *Raft) DPrintStatus() {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	entryTerm := []int{}
	for _, entry := range rf.log {
		entryTerm = append(entryTerm, entry.Term)
	}
	DPrintf("%d->{state:%d, currentTerm:%d, votedFor:%d, Next:%d, commitIndex:%d, lastApplied:%d, "+
		//"entries:%+v, " +
		"nextIndex:%+v, matchIndex:%+v}",
		rf.me, rf.state, rf.currentTerm, rf.votedFor, len(rf.log), rf.commitIndex, rf.lastApplied,
		//rf.log,
		rf.nextIndex, rf.matchIndex)
}

//
// save Raft's persistent state to stable storage,
// where it can later be retrieved after a crash and restart.
// see paper's Figure 2 for a description of what should be persistent.
//
func (rf *Raft) persist() {
	// Your code here (2C).
	// Example:
	// w := new(bytes.Buffer)
	// e := labgob.NewEncoder(w)
	// e.Encode(rf.xxx)
	// e.Encode(rf.yyy)
	// data := w.Bytes()
	// rf.persister.SaveRaftState(data)

	w := new(bytes.Buffer)
	e := labgob.NewEncoder(w)
	e.Encode(rf.currentTerm)
	e.Encode(rf.votedFor)
	e.Encode(rf.log)
	data := w.Bytes()
	rf.persister.SaveRaftState(data)
}

//
// restore previously persisted state.
//
func (rf *Raft) readPersist(data []byte) {
	if data == nil || len(data) < 1 { // bootstrap without any state?
		return
	}
	// Your code here (2C).
	// Example:
	// r := bytes.NewBuffer(data)
	// d := labgob.NewDecoder(r)
	// var xxx
	// var yyy
	// if d.Decode(&xxx) != nil ||
	//    d.Decode(&yyy) != nil {
	//   error...
	// } else {
	//   rf.xxx = xxx
	//   rf.yyy = yyy
	// }

	r := bytes.NewBuffer(data)
	d := labgob.NewDecoder(r)
	var currentTerm int
	var voteFor int
	var log []Entry
	if d.Decode(&currentTerm) != nil ||
		d.Decode(&voteFor) != nil ||
		d.Decode(&log) != nil {
		DPrintf("%d->fatal:readPersist failed", rf.me)
	} else {
		rf.currentTerm = currentTerm
		rf.votedFor = voteFor
		rf.log = log
	}
}

//
// example RequestVote RPC arguments structure.
// field names must start with capital letters!
//
type RequestVoteArgs struct {
	// Your data here (2A, 2B).
	Term         int
	CandidateId  int
	LastLogIndex int
	LastLogTerm  int
}

//
// example RequestVote RPC reply structure.
// field names must start with capital letters!
//
type RequestVoteReply struct {
	// Your data here (2A).
	Term        int
	VoteGranted bool
}

type requestVoteRequestEvent struct {
	args  *RequestVoteArgs
	reply *RequestVoteReply
	done  chan bool
}

//
// example RequestVote RPC handler.
//
func (rf *Raft) RequestVote(args *RequestVoteArgs, reply *RequestVoteReply) {
	// Your code here (2A, 2B).
	done := make(chan bool)
	rf.e <- &event{&requestVoteRequestEvent{args: args, reply: reply, done: done}}
	<-done
}

type AppendEntriesArgs struct {
	Term         int
	LeaderId     int
	PrevLogIndex int
	PrevLogTerm  int
	Entries      []Entry
	LeaderCommit int
}
type AppendEntriesReply struct {
	Term          int
	Success       bool
	ConflictIndex int
	ConflictTerm  int
}
type appendEntriesRequestEvent struct {
	args  *AppendEntriesArgs
	reply *AppendEntriesReply
	done  chan bool
}

func (rf *Raft) AppendEntries(args *AppendEntriesArgs, reply *AppendEntriesReply) {
	done := make(chan bool)
	rf.e <- &event{&appendEntriesRequestEvent{args: args, reply: reply, done: done}}
	<-done
}

//
// example code to send a RequestVote RPC to a server.
// server is the index of the target server in rf.peers[].
// expects RPC arguments in args.
// fills in *reply with RPC reply, so caller should
// pass &reply.
// the types of the args and reply passed to Call() must be
// the same as the types of the arguments declared in the
// handler function (including whether they are pointers).
//
// The labrpc package simulates a lossy network, in which servers
// may be unreachable, and in which requests and replies may be lost.
// Call() sends a request and waits for a reply. If a reply arrives
// within a timeout interval, Call() returns true; otherwise
// Call() returns false. Thus Call() may not return for a while.
// A false return can be caused by a dead server, a live server that
// can't be reached, a lost request, or a lost reply.
//
// Call() is guaranteed to return (perhaps after a delay) *except* if the
// handler function on the server side does not return.  Thus there
// is no need to implement your own timeouts around Call().
//
// look at the comments in ../labrpc/labrpc.go for more details.
//
// if you're having trouble getting RPC to work, check that you've
// capitalized all field names in structs passed over RPC, and
// that the caller passes the address of the reply struct with &, not
// the struct itself.
//
func (rf *Raft) sendRequestVote(server int, args *RequestVoteArgs, reply *RequestVoteReply) bool {
	ok := rf.peers[server].Call("Raft.RequestVote", args, reply)
	return ok
}

//
// the service using Raft (e.g. a k/v server) wants to start
// agreement on the next command to be appended to Raft's log. if this
// server isn't the leader, returns false. otherwise start the
// agreement and return immediately. there is no guarantee that this
// command will ever be committed to the Raft log, since the leader
// may fail or lose an election. even if the Raft instance has been killed,
// this function should return gracefully.
//
// the first return value is the index that the command will appear at
// if it's ever committed. the second return value is the current
// term. the third return value is true if this server believes it is
// the leader.
//
func (rf *Raft) Start(command interface{}) (int, int, bool) {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	// Your code here (2B).
	index := len(rf.log)
	term := rf.currentTerm
	isLeader := rf.state == leader

	if isLeader {
		DPrintf("%d->Start a new command", rf.me)
		rf.log = append(rf.log, Entry{Index: index, Term: term, Command: command})
		rf.persist()
	}

	return index, term, isLeader
}

//
// the tester doesn't halt goroutines created by Raft after each test,
// but it does call the Kill() method. your code can use killed() to
// check whether Kill() has been called. the use of atomic avoids the
// need for a lock.
//
// the issue is that long-running goroutines use memory and may chew
// up CPU time, perhaps causing later tests to fail and generating
// confusing debug output. any goroutine with a long-running loop
// should call killed() to check whether it should stop.
//
func (rf *Raft) Kill() {
	atomic.StoreInt32(&rf.dead, 1)
	// Your code here, if desired.
	close(rf.stopped)
}

func (rf *Raft) killed() bool {
	z := atomic.LoadInt32(&rf.dead)
	return z == 1
}

func afterBetween(min time.Duration, max time.Duration) <-chan time.Time {
	rand := rand.New(rand.NewSource(time.Now().UnixNano()))
	d, delta := min, (max - min)
	if delta > 0 {
		d += time.Duration(rand.Int63n(int64(delta)))
	}
	return time.After(d)
}

func (rf *Raft) loop() {
	DPrintf("%d->server started", rf.me)
	for !rf.killed() {
		switch rf.State() {
		case follower:
			rf.followerLoop()
		case candidate:
			rf.candidateLoop()
		case leader:
			rf.leaderLoop()
		}
	}
	DPrintf("%d->server shutdown", rf.me)
}

func (rf *Raft) followerLoop() {
	DPrintf("%d->follower loop", rf.me)

	timeoutChan := afterBetween(electionTimeout, electionTimeout*2)
	for rf.State() == follower {
		update := false

		select {
		case <-rf.stopped:
			return
		case event := <-rf.e:
			switch target := event.target.(type) {
			case *appendEntriesRequestEvent:
				update = rf.processAppendEntriesRequest(target)
			case *requestVoteRequestEvent:
				update = rf.processRequestVoteRequest(target)
			}
		case <-timeoutChan:
			rf.withLock(func() {
				rf.state = candidate
			})
		}

		if update {
			timeoutChan = afterBetween(electionTimeout, electionTimeout*2)
		}

		rf.DPrintStatus()
	}
}

func (rf *Raft) candidateLoop() {
	DPrintf("%d->candidate loop", rf.me)

	var timeoutChan <-chan time.Time
	newElection := true
	for rf.State() == candidate {
		if newElection {
			timeoutChan = afterBetween(electionTimeout, electionTimeout*2)
			go rf.startElection()
			newElection = false
		}

		select {
		case <-rf.stopped:
			return
		case event := <-rf.e:
			switch target := event.target.(type) {
			case *appendEntriesRequestEvent:
				rf.withLock(func() {
					rf.state = follower
				})
			case *upgradeToLeaderEvent:
				rf.withLock(func() {
					rf.state = leader

					length := len(rf.peers)
					rf.matchIndex = make([]int, length)
					rf.nextIndex = make([]int, length)

					nextIndex := len(rf.log)
					for i, _ := range rf.nextIndex {
						rf.nextIndex[i] = nextIndex
					}
				})
			case *requestVoteRequestEvent:
				rf.processRequestVoteRequest(target)
			case *requestVoteReplyEvent:
				rf.withLock(func() {
					if term := target.reply.Term; term > rf.currentTerm {
						rf.onReceiveLargerTerm(term)
					}
				})
			}
		case <-timeoutChan:
			newElection = true
		}

		rf.DPrintStatus()
	}
}

type upgradeToLeaderEvent struct {
}
type requestVoteReplyEvent struct {
	reply *RequestVoteReply
}

func (rf *Raft) startElection() {
	rf.mu.Lock()

	DPrintf("%d->start election", rf.me)
	rf.currentTerm += 1
	rf.votedFor = rf.me
	lastLogIndex := len(rf.log) - 1
	args := &RequestVoteArgs{
		Term:         rf.currentTerm,
		CandidateId:  rf.me,
		LastLogIndex: lastLogIndex,
		LastLogTerm:  rf.log[lastLogIndex].Term,
	}

	rf.persist()
	rf.mu.Unlock()

	peersCount := len(rf.peers)
	majority := (peersCount)/2 + 1

	voteCount := 1
	voteChannel := make(chan bool)
	for server, _ := range rf.peers {
		if server == rf.me {
			continue
		}
		go func(server int) {
			reply := &RequestVoteReply{}
			rf.sendRequestVote(server, args, reply)
			select {
			case voteChannel <- reply.VoteGranted:
			default:
			}
			rf.e <- &event{target: &requestVoteReplyEvent{reply: reply}}
		}(server)
	}

	for done := 1; voteCount < majority && done < peersCount; done++ {
		if <-voteChannel {
			voteCount += 1
		}
	}

	if voteCount >= majority {
		rf.e <- &event{target: &upgradeToLeaderEvent{}}
	}
}

func (rf *Raft) leaderLoop() {
	DPrintf("%d->leader loop", rf.me)

	heartBeat := true
	syncLog := true
	var timeoutChan <-chan time.Time
	for rf.State() == leader {
		if syncLog {
			timeoutChan = time.After(heartBeatTimeout)
			go rf.syncLog(heartBeat)
			heartBeat = false
			syncLog = false
		}

		select {
		case <-rf.stopped:
			return
		case event := <-rf.e:
			switch target := event.target.(type) {
			case *appendEntriesRequestEvent:
				rf.processAppendEntriesRequest(target)
			case *appendEntriesReplyEvent:
				rf.processAppendEntriesReply(target)
			case *requestVoteRequestEvent:
				rf.processRequestVoteRequest(target)
			}
		case <-timeoutChan:
			syncLog = true
		}

		rf.DPrintStatus()
	}
}

func (rf *Raft) processAppendEntriesReply(target *appendEntriesReplyEvent) {
	rf.mu.Lock()
	defer rf.mu.Unlock()
	server := target.server
	matchIndex := target.matchIndex
	reply := target.reply

	DPrintf("%d->AER:%d,matchIndex:%d,reply:%+v", rf.me, server, matchIndex, reply)

	if reply.Term < rf.currentTerm { // ignore latency response
		return
	}

	if reply.Term > rf.currentTerm {
		rf.onReceiveLargerTerm(reply.Term)
		return
	}

	if !reply.Success {
		nextIndex := -1
		for index, entry := range rf.log {
			if entry.Term == reply.ConflictTerm {
				nextIndex = index
				break
			}
		}
		if nextIndex == -1 {
			nextIndex = reply.ConflictIndex
		}
		rf.nextIndex[server] = min(nextIndex, len(rf.log))

		DPrintf("%d->after fail from server:%d nextIndex:%d", rf.me, server, rf.nextIndex[server])
		return
	}

	rf.nextIndex[server] = min(matchIndex+1, len(rf.log))
	rf.matchIndex[server] = matchIndex

	//rf.matchIndex[server] = max(rf.matchIndex[server], matchIndex)
	//rf.nextIndex[server] = max(rf.nextIndex[server], matchIndex+1)

	N := 0
	majority := len(rf.peers) / 2
	for _, i := range rf.matchIndex {
		count := 0
		for _, j := range rf.matchIndex {
			if j >= i {
				count += 1
			}
		}
		if count >= majority {
			N = max(N, i)
		}
	}
	DPrintf("%d->N:%d", rf.me, N)

	if rf.log[N].Term == rf.currentTerm {
		rf.commitIndex = max(rf.commitIndex, N)

		if rf.lastApplied < rf.commitIndex {
			DPrintf("%d->apply(%d:%d]", rf.me, rf.lastApplied, rf.commitIndex)
			rf.applyCommitLog()
		}
	}
}

type appendEntriesReplyEvent struct {
	server     int
	matchIndex int
	reply      *AppendEntriesReply
}

func (rf *Raft) syncLog(isHeartBeat bool) {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	for index := range rf.peers {
		server := index
		if server == rf.me {
			continue
		}
		rf.sendAppendEntriesRequest(server, isHeartBeat)
	}
}

func (rf *Raft) sendAppendEntries(server int, args *AppendEntriesArgs, reply *AppendEntriesReply) bool {
	ok := rf.peers[server].Call("Raft.AppendEntries", args, reply)
	return ok
}

func (rf *Raft) sendAppendEntriesRequest(server int, isHeartBeat bool) {
	prevLogIndex := rf.nextIndex[server] - 1
	var entries []Entry
	if !isHeartBeat && len(rf.log)-1 >= rf.nextIndex[server] {
		entries = rf.log[prevLogIndex+1:]
	}
	args := &AppendEntriesArgs{
		Term:         rf.currentTerm,
		LeaderId:     rf.me,
		PrevLogIndex: prevLogIndex,
		PrevLogTerm:  rf.log[prevLogIndex].Term,
		Entries:      entries,
		LeaderCommit: rf.commitIndex,
	}
	go func(server, matchIndex int) {
		reply := &AppendEntriesReply{}
		if ok := rf.sendAppendEntries(server, args, reply); !ok {
			return
		}
		rf.e <- &event{&appendEntriesReplyEvent{
			server:     server,
			matchIndex: matchIndex,
			reply:      reply,
		}}
	}(server, prevLogIndex+len(entries))
}
func (rf *Raft) processAppendEntriesRequest(req *appendEntriesRequestEvent) bool {
	rf.mu.Lock()
	defer rf.mu.Unlock()

	args := req.args
	reply := req.reply

	DPrintf("%d->AE:from %d,AppendEntriesArgs:"+
		"{Term:%d,LeaderId:%d,PrevLogIndex:%d,PrevLogTerm:%d,EntriesLength:%d,LeaderCommit:%d}",
		rf.me, args.LeaderId,
		args.Term, args.LeaderCommit, args.PrevLogIndex, args.PrevLogTerm, len(args.Entries), args.LeaderCommit)

	if rf.currentTerm > args.Term {
		DPrintf("%d->AE1:currentTerm > args.term", rf.me)
		reply.Term, reply.Success = rf.currentTerm, false
		close(req.done)
		return false
	}
	if rf.currentTerm < args.Term {
		DPrintf("%d->AE2:receive large term", rf.me)
		rf.onReceiveLargerTerm(args.Term)
	}

	if rf.state != follower {
		DPrintf("%d->AE:Convert to follower", rf.me)
	}
	rf.state = follower

	reply.Term = rf.currentTerm

	nextIndex := len(rf.log)
	if args.PrevLogIndex >= nextIndex {
		DPrintf("%d->AE3:does not have preLogIndex", rf.me)

		reply.ConflictIndex = nextIndex
		reply.ConflictTerm = -1
		close(req.done)
		return true
	}
	if term := rf.log[args.PrevLogIndex].Term; term != args.PrevLogTerm {
		DPrintf("%d->AE4:term not match", rf.me)
		i := args.PrevLogIndex
		for ; i > 0; i-- {
			if rf.log[i-1].Term != term {
				break
			}
		}
		reply.ConflictIndex = i
		reply.ConflictTerm = term
		close(req.done)
		return true
	}

	entriesLength := len(args.Entries)
	if args.PrevLogIndex+1 < nextIndex {
		i, j := args.PrevLogIndex+1, 0
		for i < nextIndex && j < entriesLength {
			if notEqual(&rf.log[i], &args.Entries[j]) {
				DPrintf("%d->AE5:conflict entry at %d,%d\n%+v\n%+v", rf.me, i, j, rf.log, args.Entries)
				break
			}

			i += 1
			j += 1
		}
		rf.log = rf.log[:i]
		args.Entries = args.Entries[j:]

		rf.persist()
	}

	rf.log = append(rf.log, args.Entries...)

	rf.persist()

	nextIndex = len(rf.log)
	DPrintf("%d->AE6:after AppendEntries index:%d", rf.me, nextIndex)

	if args.LeaderCommit > rf.commitIndex {
		rf.commitIndex = min(min(args.LeaderCommit, args.PrevLogIndex+entriesLength), nextIndex-1)

		DPrintf("%d->AE7:commitIndex:%d", rf.me, rf.commitIndex)
	}

	rf.applyCommitLog()

	reply.Term, reply.Success = rf.currentTerm, true

	close(req.done)
	return true
}

func notEqual(x, y *Entry) bool {
	return x.Command != y.Command || x.Term != y.Term
}

func (rf *Raft) applyCommitLog() {
	if rf.lastApplied < rf.commitIndex {
		DPrintf("%d->apply commit log [%d, %d]", rf.me, rf.lastApplied+1, rf.commitIndex)
	}

	for i := rf.lastApplied + 1; i <= rf.commitIndex; i++ {
		entry := rf.log[i]
		rf.applyCh <- ApplyMsg{
			CommandValid: true,
			CommandIndex: entry.Index,
			Command:      entry.Command,
		}
		rf.lastApplied = i
	}
}

func max(x, y int) int {
	if x > y {
		return x
	} else {
		return y
	}
}

func min(x, y int) int {
	if x < y {
		return x
	} else {
		return y
	}
}

func (rf *Raft) processRequestVoteRequest(event *requestVoteRequestEvent) bool {
	args := event.args
	reply := event.reply

	DPrintf("%d->RV:RequestVoteArgs:%+v", rf.me, args)
	rf.mu.Lock()
	defer rf.mu.Unlock()

	if rf.currentTerm > args.Term {
		reply.Term, reply.VoteGranted = rf.currentTerm, false
		close(event.done)
		return false
	}

	// term > currentTerm, update term
	if rf.currentTerm < args.Term {
		DPrintf("%d->RV:Larger Term from RequestVote", rf.me)
		rf.onReceiveLargerTerm(args.Term)
	}

	update, grantVote := false, false
	// check log is up to date
	lastLogIndex := len(rf.log) - 1
	lastLogTerm := rf.log[lastLogIndex].Term
	logIsUpToDate := args.LastLogTerm > lastLogTerm || (args.LastLogTerm == lastLogTerm && args.LastLogIndex >= lastLogIndex)
	// had voted for others
	couldVote := rf.votedFor == -1 || rf.votedFor == args.CandidateId
	if logIsUpToDate && couldVote {
		rf.votedFor = args.CandidateId
		update, grantVote = true, true
	}

	rf.persist()

	reply.Term, reply.VoteGranted = rf.currentTerm, grantVote
	DPrintf("%d->RV:RequestVoteReply:%+v to %d", rf.me, reply, args.CandidateId)

	close(event.done)
	return update
}

func (rf *Raft) onReceiveLargerTerm(term int) {
	rf.currentTerm = term
	rf.state = follower
	rf.votedFor = -1

	rf.persist()
}

//
// the service or tester wants to create a Raft server. the ports
// of all the Raft servers (including this one) are in peers[]. this
// server's port is peers[me]. all the servers' peers[] arrays
// have the same order. persister is a place for this server to
// save its persistent state, and also initially holds the most
// recent saved state, if any. applyCh is a channel on which the
// tester or service expects Raft to send ApplyMsg messages.
// Make() must return quickly, so it should start goroutines
// for any long-running work.
//
func Make(peers []*labrpc.ClientEnd, me int,
	persister *Persister, applyCh chan ApplyMsg) *Raft {
	rf := &Raft{}
	rf.peers = peers
	rf.persister = persister
	rf.me = me

	// Your initialization code here (2A, 2B, 2C).
	rf.log = []Entry{{}}
	rf.state = follower
	rf.votedFor = -1

	rf.applyCh = applyCh
	rf.e = make(chan *event)
	rf.stopped = make(chan bool)
	// initialize from state persisted before a crash
	rf.readPersist(persister.ReadRaftState())

	rf.DPrintStatus()

	go rf.loop()

	return rf
}
